import cv2
import numpy as np
import torch
from copy import deepcopy
from scipy.ndimage import distance_transform_edt

def filter_by_spatial_prior(
        segmentation_input: torch.Tensor,
        filter_rules: dict,
        filled_label: int = -1
) -> torch.Tensor:
    """
    直接根据像素的Y坐标和预设的类别空间先验规则过滤分割结果 (PyTorch Tensor 版本)。

    Args:
        segmentation_input (torch.Tensor): 输入的分割结果，形状为 (1, H, W)
        filter_rules (dict): 过滤规则字典。键是需要过滤的类别ID (int)，
                              值是一个字典，包含 'max_y_ratio' (float) 表示该类别允许的最大Y坐标比例，
                              或 'min_y_ratio' (float) 表示该类别允许的最小Y坐标比例。
                              例如: {5: {'max_y_ratio': 0.4}} 表示类别5（天空）不能出现在图像高度的40%以下。
                                    {1: {'min_y_ratio': 0.6}} 表示类别1（道路）不能出现在图像高度的60%以上。
        uncertain_label (int): 用于标记被过滤掉的像素的标签，默认为 -1。

    Returns:
        torch.Tensor: 过滤后的分割结果，形状与输入 segmentation_input 保持一致。
    """

    # 获取输入的维度信息
    if segmentation_input.ndim != 3 or segmentation_input.shape[0] != 1:
        raise ValueError("segmentation_input 必须是 (1, H, W) 形状的 Tensor。")

    # 移除 batch 维度，方便处理 H, W
    processed_tensor = segmentation_input.squeeze(0) # 形状变为 (H, W)
    image_height, image_width = processed_tensor.shape

    # 创建 Y 坐标的 Tensor (H, W)
    # y_coords = torch.arange(image_height, device=processed_tensor.device).float().unsqueeze(1).repeat(1, image_width)
    # 另一种更简洁的创建 Y 坐标矩阵的方式
    y_coords = torch.arange(image_height, dtype=torch.float32,
                            device=processed_tensor.device).unsqueeze(1)
    y_coords_matrix = y_coords.expand(-1, image_width) # 形状 (H, W)

    # 遍历过滤规则
    for class_id, rules in filter_rules.items():
        # 找到当前类别在分割图中的所有像素
        mask_class = (processed_tensor == class_id)

        if 'max_y_ratio' in rules:
            max_allowed_y = int(image_height * rules['max_y_ratio'])
            # 如果像素是当前类别，但Y坐标大于允许的最大Y坐标
            abnormal_pixels_mask = mask_class & (y_coords_matrix > max_allowed_y)
            processed_tensor[abnormal_pixels_mask] = filled_label

        if 'min_y_ratio' in rules:
            min_allowed_y = int(image_height * rules['min_y_ratio'])
            # 如果像素是当前类别，但Y坐标小于允许的最小Y坐标
            abnormal_pixels_mask = mask_class & (y_coords_matrix < min_allowed_y)
            processed_tensor[abnormal_pixels_mask] = filled_label

    # 恢复原始的 (1, H, W) 形状
    return processed_tensor.unsqueeze(0)
def post_process_roads_morphology(
        segmentation_input,
        road_labels,
        unknown_label=-1,
        dila_kernel_size=(9, 9),
        open_kernel_size=(5, 5),
        min_component_size=100,
        expect_labels_lst=None,
        fill_unknown=False
):
    import time
    t0 = time.time()
    # 判断类型 & 记录属性
    is_tensor = isinstance(segmentation_input, torch.Tensor)
    orig_dev = segmentation_input.device
    orig_dtype = segmentation_input.dtype

    if is_tensor:
        seg_map = segmentation_input.squeeze().cpu().numpy()
    elif isinstance(segmentation_input, np.ndarray):
        # 直接转成 int32，就能放下负值
        seg_map = segmentation_input
    else:
        raise TypeError("segmentation_input 必须是 Tensor 或 ndarray")
    t1 = time.time()
    print(f"\t\t转tensor: {round((t1 - t0) * 1000, 3)} ms")

    processed = seg_map.copy()

    # 1. 非道路去噪 - 直接基于连通区域大小过滤
    non_road_mask = np.isin(processed, road_labels + [unknown_label],
                            invert=True).astype(np.uint8) * 255
    # 直接进行连通域分析
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(
        non_road_mask, connectivity=8)

    # 创建过滤后的非道路掩码
    filtered_non_road_mask = np.zeros_like(non_road_mask)
    for i in range(1, num_labels):  # 跳过背景(i=0)
        area = stats[i, cv2.CC_STAT_AREA]
        if area >= min_component_size:  # 保留足够大的区域
            filtered_non_road_mask[labels == i] = 255

    # 将小组件置为unknown
    small_components = (non_road_mask == 255) & (filtered_non_road_mask == 0)
    processed[small_components] = unknown_label

    t2 = time.time()
    print(f"\t\t去噪: {round((t2 - t1) * 1000, 3)} ms")

    # # 2. 道路膨胀合并
    # dila_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, dila_kernel_size)
    # # 创建所有道路的合并掩码
    # all_roads_mask = np.zeros_like(seg_map, dtype=np.uint8)
    # for road_label in road_labels:
    #     all_roads_mask[seg_map == road_label] = road_label
    # # 一次性膨胀所有道路
    # dilated_roads = cv2.dilate(all_roads_mask.astype(np.uint8), dila_kernel,
    #                            iterations=1)
    # # 只在安全区域应用膨胀
    # safe = ~np.isin(seg_map, expect_labels_lst)
    # mask_to_update = (dilated_roads > 0) & (all_roads_mask == 0) & safe
    # processed[mask_to_update] = dilated_roads[mask_to_update]
    #
    t3 = time.time()
    # print(f"\t\t道路膨胀: {round((t3 - t2) * 1000, 3)} ms")

    # 3. 填充 unknown（可选）
    if fill_unknown:
        unk_mask = (processed == unknown_label)
        if np.any(unk_mask):
            dist, inds = distance_transform_edt(
                unk_mask,
                return_distances=True,
                return_indices=True
            )
            rr, cc = inds
            processed[unk_mask] = processed[rr[unk_mask], cc[unk_mask]]

    t4 = time.time()
    print(f"\t\t填充: {round((t4 - t3) * 1000, 3)} ms")

    # 4. 返回同类型输出
    if is_tensor:
        out_np = processed.astype(seg_map.dtype)
        out = torch.from_numpy(out_np).to(device=orig_dev, dtype=orig_dtype)
        # 保持 batch 维度
        if segmentation_input.ndim == 4 or (segmentation_input.ndim == 3 and segmentation_input.shape[0] > 1):
            out = out.unsqueeze(0)

        return out
    else:
        return processed


allowed_labels_tensor = None
def fill_uncertain_pixels_temporal(
        current_segmentation_input,
        previous_segmentation_map=None,  # New parameter for temporal information
        uncertain_label=-1,
        default_fill_value=-1,
        allowed_previous_labels: list = None
):
    """
    使用时间信息（前一帧）和欧几里得距离变换填充分割图中类别为 uncertain_label 的像素。
    优先使用前一帧的对应像素标签进行填充。如果前一帧信息不可用或本身不确定，则回退到距离变换填充。

    Args:
        current_segmentation_input (torch.Tensor or np.ndarray): 当前帧的分割结果。
                                                            如果是Tensor，会自动转换为NumPy处理，
                                                            处理结束后再转回Tensor。
                                                            每个像素包含类别标签 (0..N-1) 或 uncertain_label。
        previous_segmentation_map (torch.Tensor or np.ndarray, optional): 前一帧的分割结果。
                                                                       如果提供，将用于时间填充。
                                                                       默认为 None。
        uncertain_label (int): 表示不确定类别的标签，默认为 -1。
        default_fill_value (int): 如果不确定像素周围没有确定类别（即它是孤立的），
                                  则填充的默认值，默认为 0 (背景)。

    Returns:
        torch.Tensor or np.ndarray: 填充后的分割结果。
                                     如果输入是Tensor，返回Tensor；否则返回NumPy数组。
    """
    global allowed_labels_tensor
    if (current_segmentation_input.ndim != 3 or
            current_segmentation_input.shape[0] != 1):
        raise ValueError("segmentation_input 必须是 (1, H, W) 形状的 Tensor。")

    filled_map = current_segmentation_input.clone()
    device = filled_map.device

    if allowed_previous_labels is not None and allowed_labels_tensor is None:
        allowed_labels_tensor = torch.tensor(allowed_previous_labels,
                                             device=device)

    # 找到所有不确定像素的位置
    uncertain_mask = (filled_map == uncertain_label)

    # --- 优先进行时间填充 ---
    if previous_segmentation_map is not None:

        # 1. 找到当前帧中不确定像素的位置
        # 2. 找到前一帧中对应位置不是不确定标签的像素
        # 3. 找到前一帧中对应位置的标签在允许列表中的像素
        temporal_fill_mask = (uncertain_mask & \
                                 (previous_segmentation_map != uncertain_label) & \
                                 torch.isin(previous_segmentation_map, allowed_labels_tensor)
                              )
        # 使用前一帧的值填充这些像素
        filled_map = torch.where(temporal_fill_mask, previous_segmentation_map,
                                 filled_map)

        # --- 填充剩余的不确定类别 ---
        # 再次找到所有不确定像素（那些未被时间信息填充的）
        remaining_uncertain_mask = (filled_map == uncertain_label)

        # 将这些剩余的不确定像素设置为默认值
        filled_map[remaining_uncertain_mask] = default_fill_value

        return filled_map

    else:
        return current_segmentation_input

